# 导入包
import pandas as pd
import statsmodels.formula.api as smf
from sklearn.model_selection import train_test_split
from sqlalchemy import create_engine
from sklearn.metrics import r2_score, mean_squared_error
from matplotlib import pyplot as plt
import statsmodels.api as sm
import numpy as np
import seaborn as sns
import pymysql
db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': '123456',
    'database': 'tushare',
    'port': 3306,          # 明确指定端口
    'charset': 'utf8mb4'   # 添加字符集设置
}
engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
)
conn = pymysql.connect(**db_config)
chunk_size = 10000

# 获取华夏银行日线数据
df = pd.read_sql_query(
        """
        SELECT d.*, m.net_mf_vol , m.sell_elg_vol , m.buy_elg_vol, m.sell_lg_vol, m.buy_lg_vol , i.closes as i_closes, i.vol as i_vol
        FROM date_1 d
        JOIN moneyflows m 
        ON d.ts_code = m.ts_code AND d.trade_date = m.trade_date
        LEFT JOIN index_daily i ON  d.trade_date = i.trade_date and i.ts_code='000001.SH'
        WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' AND d.ts_code='002229.SZ'
        """, 
        conn, 
        chunksize=chunk_size
    )
df1 = pd.concat(df, ignore_index=True)

df1.head()
# 增加一列，股票的涨跌幅
df1['zd_close'] = round((df1['closes']-df1['closes'].shift(1))/df1['closes'].shift(1),2)
df1['hz_close'] = round((df1['i_closes']-df1['i_closes'].shift(1))/df1['i_closes'].shift(1),2)
df1['hz_vol'] = round((df1['i_vol'].shift(1)-df1['i_vol'].shift(2))/df1['i_vol'].shift(2),2)
# 处理缺失值
df1 = df1.dropna(subset=['zd_close', 'hz_close', 'hz_vol'])


numeric_cols = df1.select_dtypes(include=['number']).columns.tolist()
# 筛选合适的自变量
excluded = ['zd_close', 'ts_code', 'trade_date', 'id', 'pre_closes', 'changes', 'pct_chg', 'opens', 'high', 'low', 'closes', 
            ]  # 根据实际列名调整
predictors = [col for col in numeric_cols if col not in excluded]

# 排除方差为零的列 移除所有值相同的列，避免共线性问题
# predictors = [col for col in predictors if df1[col].nunique() > 1]

# 检查样本量与自变量数量
n = len(df1)
k = len(predictors)
if n <= k + 1:
    print(f"错误：样本数{n}不足，至少需要{k + 2}个样本。")
else:
    formula = 'zd_close ~ ' + ' + '.join(predictors)
    results1 = smf.ols(formula, data=df1).fit()
    print(results1.summary())

plt.figure(figsize=(10, 6))
plt.scatter(results1.fittedvalues, df1['zd_close'])
plt.title('Fitted Values vs. Observed Values')
plt.xlabel('Fitted Values')
plt.ylabel('Observed Values')
plt.show()


features = df1[['vol', 'amount']]
correlation_matrix = features.corr()

# Plotting a heatmap for the correlation matrix
plt.figure(figsize=(10, 6))
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', cbar=True, fmt='.2f', linewidths=.5)
plt.title('Correlation Matrix between Features')
plt.show()
